<script setup lang="ts">
import type { Color } from "three";
import { Vector2, Vector3 } from "three";
import { Line2 } from "three/examples/jsm/lines/Line2";
import { LineMaterial } from "three/examples/jsm/lines/LineMaterial";
import { LineGeometry } from "three/examples/jsm/lines/LineGeometry";
import { computed, onUnmounted, watch } from "vue";
import type { TresColor } from "@tresjs/core";
import { normalizeColor, useTresContext } from "@tresjs/core";

type Points = (
  | Vector3
  | Vector2
  | [number, number, number]
  | [number, number]
  | number
)[];
type VertexColors = Array<TresColor>;
export interface LineProps {
  points: Points;
  vertexColors?: VertexColors | null;
  color?: TresColor;
  lineWidth?: number;
  worldUnits?: boolean;
  alphaToCoverage?: boolean;
  dashed?: boolean;
  dashSize?: number;
  gapSize?: number;
  dashScale?: number;
  dashOffset?: number;
  opacity?: number;
}

const props = withDefaults(defineProps<LineProps>(), {
  vertexColors: null,
  color: "white",
  lineWidth: 1,
  worldUnits: false,
  alphaToCoverage: false,
  dashed: false,
  dashSize: 1,
  gapSize: 1,
  dashScale: 1,
  dashOffset: 0,
  opacity: 1,
});

type PropsType = typeof props;

function getInterpolatedVertexColors(
  vertexColors: VertexColors | null,
  numPoints: number
): Color[] {
  if (!vertexColors || vertexColors.length === 0) {
    return Array.from({ length: numPoints }).fill(
      normalizeColor(props.color)
    ) as Color[];
  }
  if (vertexColors.length === 1) {
    return Array.from({ length: numPoints }).fill(
      normalizeColor(vertexColors[0])
    ) as Color[];
  }
  if (vertexColors.length === numPoints) {
    return vertexColors.map(normalizeColor);
  }

  const numSegments = numPoints - 1;
  const mappedColors = vertexColors.map(normalizeColor);
  if (closed) {
    mappedColors.push(mappedColors[0].clone());
  }

  const iColors: Color[] = [mappedColors[0]];
  const divisions = numSegments / (mappedColors.length - 1);
  for (let i = 1; i < numSegments; i++) {
    const alpha = (i % divisions) / divisions;
    const colorIndex = Math.floor(i / divisions);
    iColors.push(
      mappedColors[colorIndex].clone().lerp(mappedColors[colorIndex + 1], alpha)
    );
  }
  iColors.push(mappedColors[mappedColors.length - 1]);

  return iColors;
}

const lineMaterial = new LineMaterial();
const lineGeometry = new LineGeometry();
const line = new Line2(lineGeometry, lineMaterial);
const sizes = useTresContext().sizes;
const hasVertexColors = computed(() => Array.isArray(props.vertexColors));

function updateLineMaterial(material: LineMaterial, props: PropsType) {
  material.color = normalizeColor(props.color);
  material.linewidth = props.lineWidth;
  material.alphaToCoverage = props.alphaToCoverage;
  material.worldUnits = props.worldUnits;
  material.vertexColors = Array.isArray(props.vertexColors);
  material.dashed = props.dashed;
  material.dashScale = props.dashScale;
  material.dashSize = props.dashSize;
  material.dashOffset = props.dashOffset;
  material.gapSize = props.gapSize;
  material.needsUpdate = true;
  material.opacity = props.opacity;
}

function updateLineGeometry(
  geometry: LineGeometry,
  points: Points,
  vertexColors: VertexColors | null
) {
  const pValues = points
    .map((p) => {
      if (p instanceof Vector3) {
        return [p.x, p.y, p.z];
      } else if (p instanceof Vector2) {
        return [p.x, p.y, 0];
      } else if (Array.isArray(p) && p.length === 2) {
        return [p[0], p[1], 0];
      } else {
        return p;
      }
    })
    .flat();
  geometry.setPositions(pValues.flat());

  const colors = getInterpolatedVertexColors(vertexColors, points.length)
    .map((c) => c.toArray())
    .flat();
  geometry.setColors(colors);

  line.computeLineDistances();
}

updateLineMaterial(lineMaterial, props);
updateLineGeometry(lineGeometry, props.points, props.vertexColors);
line.computeLineDistances();

watch(
  () => [
    props.color,
    props.lineWidth,
    props.alphaToCoverage,
    props.worldUnits,
    hasVertexColors,
    props.dashed,
    props.dashScale,
    props.dashSize,
    props.dashOffset,
  ],
  () => updateLineMaterial(lineMaterial, props)
);
watch([props.points, props.vertexColors], () =>
  updateLineGeometry(lineGeometry, props.points, props.vertexColors)
);
watch(
  () => props.vertexColors,
  () => updateLineGeometry(lineGeometry, props.points, props.vertexColors)
);
watch(
  () => props.points,
  () => updateLineGeometry(lineGeometry, props.points, props.vertexColors)
);
watch(
  [sizes.height, sizes.width],
  () =>
    (lineMaterial.resolution = new Vector2(
      sizes.width.value,
      sizes.height.value
    ))
);

onUnmounted(() => {
  lineGeometry.dispose();
  lineMaterial.dispose();
});
</script>

<template>
  <primitive :object="line" />
</template>
